package main

import (
	"context"
	"net"
	"os"
	"time"

	"github.com/sirupsen/logrus"
)

type DnsResolver struct {
}

func NewDnsResolver() *DnsResolver {
	return &DnsResolver{}
}

// 自定义DNS解析器
func customResolver(dnsServer string) *net.Resolver {
	return &net.Resolver{
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			// 忽略传入的address，使用指定的DNS服务器
			d := net.Dialer{
				Timeout: 5 * time.Second,
			}
			return d.DialContext(ctx, "udp", dnsServer+":53")
		},
	}
}

func (r *DnsResolver) LookupIPAddr(domain string, dns string) {
	//domain := "test.com"
	var dnsServer string

	dnsServer = dns

	logrus.Infof("使用DNS服务器 %s 查询 %s...", dnsServer, domain)
	resolver := customResolver(dnsServer)
	ips, err := resolver.LookupIPAddr(context.Background(), domain)
	if err != nil {
		logrus.Infof("查询失败: %v", err)
		return
	}
	for _, ip := range ips {
		logrus.Infof("解析结果: %s", ip.IP)
	}
}

func main() {
	if len(os.Args) < 3 {
		logrus.Infof("使用样例: dns-client face13333.com 8.8.8.8")
		return
	}
	domainName := os.Args[1]
	server := os.Args[2]
	logrus.Infof("domainName=%v,server=%v", domainName, server)
	dnsResolver := NewDnsResolver()
	dnsResolver.LookupIPAddr(domainName, server)
}
